# 导入所需的模块和库
import copy
import inspect
from pprint import pformat
import traceback
from types import GenericAlias
from typing import get_origin, Annotated
import subprocess

# 用于存储已注册工具的函数和其描述信息
_TOOL_HOOKS = {}
_TOOL_DESCRIPTIONS = {}


# 装饰器函数，用于注册工具
def register_tool(func: callable):
    # 获取工具的名称、描述和参数信息
    tool_name = func.__name__
    tool_description = inspect.getdoc(func).strip()
    python_params = inspect.signature(func).parameters
    tool_params = []

    # 遍历工具的参数
    for name, param in python_params.items():
        annotation = param.annotation
        # 检查是否有类型注解
        if annotation is inspect.Parameter.empty:
            raise TypeError(f"参数 `{name}` 缺少类型注解")
        # 检查类型注解是否为 Annotated
        if get_origin(annotation) != Annotated:
            raise TypeError(f"`{name}` 的注解类型必须为 typing.Annotated")

        # 获取注解的类型、描述和是否必需信息
        typ, (description, required) = annotation.__origin__, annotation.__metadata__
        typ: str = str(typ) if isinstance(typ, GenericAlias) else typ.__name__
        # 检查描述和是否必需信息的类型
        if not isinstance(description, str):
            raise TypeError(f"`{name}` 的描述必须为字符串")
        if not isinstance(required, bool):
            raise TypeError(f"`{name}` 的必需属性必须为布尔值")

        # 将参数信息添加到工具参数列表中
        tool_params.append({
            "name": name,
            "description": description,
            "type": typ,
            "required": required
        })

    # 构建工具的定义信息
    tool_def = {
        "name": tool_name,
        "description": tool_description,
        "params": tool_params
    }

    # 打印注册的工具信息
    print("[已注册工具] " + pformat(tool_def))

    # 将工具函数和描述信息添加到全局变量中
    _TOOL_HOOKS[tool_name] = func
    _TOOL_DESCRIPTIONS[tool_name] = tool_def

    return func


# 调用注册的工具并返回结果
def dispatch_tool(tool_name: str, tool_params: dict) -> str:
    if tool_name not in _TOOL_HOOKS:
        return f"工具 `{tool_name}` 未找到。请使用已提供的工具。"
    tool_call = _TOOL_HOOKS[tool_name]
    try:
        ret = tool_call(**tool_params)
    except:
        # 处理工具调用过程中的异常
        ret = traceback.format_exc()
    return str(ret)


# 获取已注册工具的描述信息
def get_tools() -> dict:
    return copy.deepcopy(_TOOL_DESCRIPTIONS)


# ==========================自定义工具=============================

# 注册一个在Linux shell中运行命令的工具
@register_tool
def get_shell(
        query: Annotated[str, '应在Linux shell中运行的命令', True],
) -> str:
    """
       使用shell运行命令
    """
    if not isinstance(query, str):
        raise TypeError("命令必须是字符串")
    try:
        result = subprocess.run(query, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                text=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        return e.stderr


# 注册一个查询数据库的工具
@register_tool
def query_db(
        database: Annotated[str, '要查询的数据库(Milvus、Mariadb、MySQL)', True],
        sql: Annotated[str, '数据库查询语句', True],
) -> str:
    """
       Query Database by `db_type` type
    """
    match database.lower():
        case "mariadb":
            result = "Mariadb --> " + sql
        case "mysql":
            result = "MySQL --> " + sql
        case "milvus":
            result = "Milvus --> " + sql
        case _:
            result = f"不支持的数据库类型: {database}"
            # raise ValueError(f"不支持的数据库: {database}")

    # 返回数据库
    return result


if __name__ == "__main__":
    # print(dispatch_tool("get_shell", {"query": "pwd"}))
    print(get_tools())
    print(query_db("Mysql", "SELECT * FROM table1 WHERE 1==1"))
